#!/usr/bin/env python
#Albert Myers
#Python 2.6.6 script to decode binary traces
#Works with mcfTrace, mlsTrace, mcfTRaptor, and mlvCFiat

import sys
import struct
import binascii
def extenPrecEndian( temp ):
	hex_string =  "".join("%02x" % b for b in bytearray(temp[::-1]))
	return "0x" + hex_string

def readmcfTrace(i, o):
	# tid ( 1 byte ), addr (8 bytes), target addr (8 bytes) , branch type (1 byte), conditional/uncon (1 byte), dir/indir (1 byte), outcome (1 byte), Stream Length (4 byte)
    tid = i.read(1)
    while tid:
			#print tid
        	addr = i.read(8)
        	target = i.read(8)
        	type = i.read(1)

		o.write( str( ord(tid) ) + ", "  )
		o.write( '0x%016x' % (struct.unpack('Q', addr)[0]) + ", " )
		o.write( '0x%016x' % (struct.unpack('Q', target)[0]) + ", " )
		type = ord(type)
		#print type
		if type == 0:
			o.write( "U, I, T, " )
		elif type == 1:
			o.write( "U, D, T, " )
		elif type == 2:
			o.write( "C, D, T, " )
		elif type == 3:
			o.write( "C, D, NT, " )
		else:
			print("Error decoding mcfTrace\n")
			exit(1)
		SL = i.read(4)
		o.write('%01i' % (struct.unpack('i', SL)[0])) 
		o.write("\n")
		tid = i.read(1)

def readmlsTrace(i, o):
	# tid ( 1 byte ), ins addr (8 bytes), op addr (8 bytes), op size (1 byte), value (op size bytes) 
        meta = ord(i.read(1))
	if meta == 0:
		tid = i.read(1)
		while tid:
			insaddr = i.read(8)
			opaddr = i.read(8)
			opsize = ord(i.read(1))
			value = i.read(opsize)	
			o.write( str( ord(tid) ) + ", "  )
			o.write( '0x%016x' % (struct.unpack('Q', insaddr)[0]) + ", " )
			o.write( '0x%016x' % (struct.unpack('Q', opaddr)[0]) + ", " )
			o.write( str( opsize ) + ", " )
			format = '0x%0' + str(opsize*2) + "x"
			if opsize == 1:
				num = ord(value)
				o.write( format % num + '\n' )
			elif opsize == 2:
				num = struct.unpack('<H', value)[0]
				o.write( format % num + '\n' )
			elif opsize == 4:
				num = struct.unpack('<I', value)[0]
				o.write( format % num + '\n' )
			elif opsize == 8:
				num = struct.unpack('<Q', value)[0]
				o.write( format % num + '\n' )
			elif opsize == 10:
				# don't think python will support extended prec, just reverse endian and print as string
				o.write(extenPrecEndian(value) + '\n')
			else:
				# anything else we can't say a lot about without seeing the instruction (i.e. sse load could be packed for fp or double)
				# just reverse endianess for whole string and print
				o.write( "0x" + "{0:0>{width}}".format( binascii.hexlify(value), width = opsize ) + '\n' )	
			
			tid = i.read(1)
	# tid ( 1 byte), type (1 byte), addr (8 bytes), op addr (8 bytes), op size (1 byte), value (op size bytes)
	else:
		tid = i.read(1)
		while tid:
			desctype = ord(i.read(1))
			insaddr = i.read(8)
			opaddr = i.read(8)
			opsize = ord(i.read(1))
			value = i.read(opsize)	
			o.write( str( ord(tid) ) + ", "  )
			if desctype == 0:
				o.write("L, ")
			else:
				o.write("S, ")
			o.write( '0x%016x' % (struct.unpack('Q', insaddr)[0]) + ", " )
			o.write( '0x%016x' % (struct.unpack('Q', opaddr)[0]) + ", " )
			o.write( str( opsize ) + ", " )
			format = '0x%0' + str(opsize*2) + "x"
			if opsize == 1:
				num = ord(value)
				o.write( format % num + '\n' )
			elif opsize == 2:
				num = struct.unpack('<H', value)[0]
				o.write( format % num + '\n' )
			elif opsize == 4:
				num = struct.unpack('<I', value)[0]
				o.write( format % num + '\n' )
			elif opsize == 8:
				num = struct.unpack('<Q', value)[0]
				o.write( format % num + '\n' )
			elif opsize == 10:
				# don't think python will support extended prec, just reverse endian and print as string
				o.write(extenPrecEndian(value) + '\n')
			else:
				# anything else we can't say a lot about without seeing the instruction (i.e. sse load could be packed for fp or double)
				# just reverse endianess for whole string and print
				o.write( "0x" + "{0:0>{width}}".format( binascii.hexlify(value), width = opsize ) + '\n' )	

			
			tid = i.read(1)
				

	
def readmcfTRaptor(i, o):
	tid = i.read(1)
	while tid:
		secondfield = struct.unpack('I', i.read(4))
		#exception - TID, 0, iCnt, Target Addr
		if secondfield == 0:
			icnt = struct.unpack('I', i.read(4))
			targetAddress = struct.unpack('Q', i.read(8))
			#write exception trace
			o.write( str( ord(tid) ) + ", "  )
			o.write( "0, ")
			o.write(  str(icnt) + ", " )
                        o.write( '0x%016x' % targetAddress + '\n' )
			tid = i.read(1)
		#mispredicted target or outcome
		else:
			o.write( str( ord(tid) ) + ", "  )
			bcnt = secondfield
			print bcnt[0]
			#o1.write(str(ord(tid)  )+ ", ")			
			#o1.write(str(bcnt[0]))
			#o1.write('\n')
			o.write(  str(bcnt[0]) )
			thirdfield = i.read(1)
			if thirdfield:
				#is mispredicted target
        	                if ord(thirdfield) == 255:
					targetAddress = struct.unpack('Q', i.read(8))
					o.write(", T, ")
                        		o.write( '0x%016x' % targetAddress + '\n' )
					tid = i.read(1)
				#is mispredicted outcome
				else:
					o.write('\n')
					tid = thirdfield
			else:
				o.write('\n')
				break
				



tool = sys.argv[1]
filename = sys.argv[2]
#o1 = open(filename+ ".bcnt.tid" + ".txt", "w")
with open(filename,"rb") as i:
	with open(filename + ".txt", "w") as o:
		

			if tool == "mcfTrace":
						readmcfTrace(i, o);        
			elif tool == "mlsTrace":
				readmlsTrace(i, o)
			elif tool == "mcfTRaptor":
				readmcfTRaptor(i, o)
			elif tool == "mlvCFiat":
				print(3)
			else:
				print("Usage: readBinaryTrace.py <Pin Tool> <Trace File>")


